import numpy as np
from scipy.optimize import fsolve
import plotly.graph_objects as go
from ipywidgets import interactive
import ipywidgets as widgets
from IPython.display import display

# Create output widget to hold the plot
output = widgets.Output()

# Display the output widget once
display(output)

def plot_with_update(**kwargs):
    # Clear previous output
    output.clear_output(wait=False)
    
    # Display new plot inside the output widget
    with output:
        plot_function(**kwargs)


def plot_function(theta, z, Pp_Pr, alpha):
    def equations(vars):
        x, y = vars
        eq1 = ((theta - 1) / theta) * x + (z / theta) - y
        eq2 = y * (1 + Pp_Pr * (1-y) * (alpha - 1)) / (alpha - (alpha - 1) * y) - x
        return [eq1, eq2]

    x_set, y_set = fsolve(equations, (0.15, 0.85))

    N = 500  # Number of points for plotting

    # Operating Line
    x_op = np.linspace(0, 1, N)
    y_op = ((theta - 1) / theta) * x_op + (z / theta)

    # Rate Transfer Line
    y_rt = np.linspace(0, 1, N)
    x_rt = y_rt * (1 + Pp_Pr * (1-y_rt) * (alpha - 1)) / (alpha - (alpha - 1) * y_rt)

    fig = go.Figure()

    fig.add_trace(go.Scatter(x=x_op, y=y_op, mode='lines', name='Operating Line', line=dict(color='red')))
    fig.add_trace(go.Scatter(x=x_rt, y=y_rt, mode='lines', name='Rate Transfer Line', line=dict(color='blue')))
    fig.add_trace(go.Scatter(x=[x_set], y=[y_set], mode='markers', name='Operating Point', 
                             marker=dict(color='lime', size=14)))

    fig.update_layout(
        title=f'Operating and Rate Transfer Equations<br>θ={theta:.3f}, z={z:.2f}, Pp/Pr={Pp_Pr:.2f}, α={alpha:.2f}',
        xaxis_title='Molar Fraction in the Retentate',
        yaxis_title='Molar Fraction in the Permeate',
        xaxis=dict(range=[0, 1]),
        yaxis=dict(range=[0, 1]),
        width=800,
        height=600
    )

    #print(f"Operating point -> x = {x_set:.4f}, y = {y_set:.4f}")
    
    fig.show()

# Create sliders
theta_slider = widgets.FloatSlider(value=0.5, min=0.001, max=0.99, step=0.001, description='$\\theta$ ')
z_slider = widgets.FloatSlider(value=0.5, min=0, max=1, step=0.01, description='$z$ ')
Pp_Pr_slider = widgets.FloatSlider(value=0.1, min=0.001, max=1, step=0.001, description='$P_p/P_r$ ')
alpha_slider = widgets.FloatSlider(value=10, min=1, max=200, step=1, description='$\\alpha$ ')

# Create and display interactive plot
# Create interactive widget with the update function
interactive_plot = interactive(plot_with_update, 
                             theta=theta_slider, 
                             z=z_slider, 
                             Pp_Pr=Pp_Pr_slider, 
                             alpha=alpha_slider)
display(interactive_plot)